#pragma once

#include <chrono>
#include <functional>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>

#include "asycipc/backpressure.hpp"
#include "asycipc/frame.hpp"
#include "asycipc/handshake.hpp"
#include "asycipc/heartbeat.hpp"
#include "asycipc/session_context.hpp"
#include "asycipc/sliding_window.hpp"
#include "asycipc/telemetry.hpp"

namespace asycipc {

struct ServerOptions {
  std::string endpoint;
  std::set<std::string> allowed_app_ids;
  std::size_t window_size = 256;
  std::chrono::milliseconds retry_base{500};
  uint32_t max_retry_attempts = 5;
  std::chrono::milliseconds heartbeat_interval{5000};
  uint32_t heartbeat_allowed_miss = 3;
  bool reliability_enabled = true;
};

struct ClientOptions {
  std::string endpoint;
  std::string app_id;
  std::vector<std::string> capabilities;
};

class Channel;

using MessageHandler =
    std::function<std::vector<Frame>(Channel&, const Frame&)>;
using ErrorHandler = std::function<void(const std::string&)>;
using BackpressureHandler = std::function<void(bool)>;

class PipeServer {
 public:
  class Builder {
   public:
    Builder& Endpoint(std::string endpoint);
    Builder& AllowApp(std::string app_id);
    Builder& OnMessage(MessageHandler handler);
    Builder& OnError(ErrorHandler handler);
    Builder& OnBackpressure(BackpressureHandler handler);
    Builder& ReliabilityEnabled(bool enabled);

    PipeServer Build() const;

   private:
    ServerOptions options_;
    MessageHandler message_handler_;
    ErrorHandler error_handler_;
    BackpressureHandler backpressure_handler_;
  };

  PipeServer(ServerOptions options, MessageHandler msg_handler,
             ErrorHandler err_handler,
             BackpressureHandler backpressure_handler);

  void Start();
  HandshakeResponse Accept(const HandshakeRequest& request,
                           SessionContext& session) const;
  std::unique_ptr<Channel> CreateChannel(SessionContext session) const;

  TelemetryCenter& telemetry() noexcept { return telemetry_; }
  const ServerOptions& options() const noexcept { return options_; }

  std::vector<Frame> ProcessInbound(Channel& channel, const std::vector<uint8_t>& raw_payload);

  std::vector<Frame> CollectMaintenanceFrames(Channel& channel);

  void ReportError(const std::string& message);
 private:
  ServerOptions options_;
  MessageHandler message_handler_;
  ErrorHandler error_handler_;
  BackpressureHandler backpressure_handler_;
  HandshakeConfig handshake_config_;
  std::shared_ptr<HandshakeManager> handshake_manager_;
  TelemetryCenter telemetry_;
};

class PipeClient {
 public:
  class Builder {
   public:
    Builder& Endpoint(std::string endpoint);
    Builder& AppId(std::string app_id);
    Builder& Capabilities(std::vector<std::string> caps);

    PipeClient Build() const;

   private:
    ClientOptions options_;
  };

  explicit PipeClient(ClientOptions options);

  HandshakeRequest BuildHandshakeRequest() const;
  const ClientOptions& options() const noexcept { return options_; }

 private:
  ClientOptions options_;
};

class Channel {
 public:
  Channel(SessionContext session, std::shared_ptr<HandshakeManager> manager);

  struct InboundFrameResult {
    std::optional<Frame> application_frame;
    std::vector<Frame> control_frames;
  };

  SessionContext& session() { return session_; }
  const SessionContext& session() const { return session_; }

  void SetMessageHandler(MessageHandler handler);
  void SetErrorHandler(ErrorHandler handler);

  void SetBackpressureHandler(BackpressureHandler handler);
  void NotifyBackpressure(bool active) const;

  void SetReliabilityModules(std::shared_ptr<SlidingWindow> window,
                             std::shared_ptr<HeartbeatMonitor> heartbeat,
                             std::shared_ptr<BackpressureController> backpressure,
                             TelemetryCenter* telemetry);
  void SetTimeProvider(
      std::function<std::chrono::steady_clock::time_point()> provider);

  InboundFrameResult HandleInboundFrame(const Frame& frame);

  std::vector<Frame> MakeDataFrames(const std::vector<uint8_t>& payload,
                                    const std::vector<uint8_t>& extensions = {});
  void OnAck(uint64_t seq);
  void OnNack(uint64_t seq);
  std::vector<SlidingWindow::RetryItem> PollRetries();
  bool ShouldSendHeartbeat() const;
  void MarkHeartbeatSent();
  void MarkHeartbeatReceived();
  bool HeartbeatExpired() const;

 private:
  std::chrono::steady_clock::time_point Now() const;

  SessionContext session_;
  std::shared_ptr<HandshakeManager> handshake_manager_;
  BackpressureHandler backpressure_handler_;
  MessageHandler message_handler_;
  ErrorHandler error_handler_;
  std::shared_ptr<SlidingWindow> sliding_window_;
  std::shared_ptr<HeartbeatMonitor> heartbeat_monitor_;
  std::shared_ptr<BackpressureController> backpressure_controller_ = nullptr;
  TelemetryCenter* telemetry_ = nullptr;
  std::function<std::chrono::steady_clock::time_point()> time_provider_;

  struct FragmentState {
    uint16_t total_chunks = 0;
    std::vector<std::vector<uint8_t>> parts;
    std::vector<bool> received;
    std::size_t received_count = 0;
    std::vector<uint8_t> extensions;
    uint8_t base_flags = 0;
    uint8_t reserved = 0;
    FrameType type = FrameType::kData;
  };

  std::unordered_map<uint64_t, FragmentState> inbound_fragments_;
};

}  // namespace asycipc
